"""
DESCRIPTION
Lists the strongest association, by type, for each target on the list. The final score,
association index, is calculated as the sum of all contributions. 

ARGUMENTS
1. input association file
2. output file with association indices

CONTACT
timo.p.erkkila@tut.fi
http://www.cs.tut.fi/~erkkila2/
"""


import csv,sys,re
from operator import itemgetter

assert sys.argv[1] != sys.argv[2]

listreader = csv.reader(open(sys.argv[1],'r'),delimiter='\t')
listwriter = csv.writer(open(sys.argv[2],'w'),delimiter='\t')

strengths = {}

featureType_pos = 1

allFeatureTypes = set()

for line in listreader:
    target,feature,pvalue,importance,correlation = line

    target_elems = re.split(':',target)
    feature_elems = re.split(':',feature)

    featureType = feature_elems[featureType_pos]

    allFeatureTypes.add(featureType)


    if strengths.get(target,'NA') == 'NA':
        strengths[target] = {}

    if strengths[target].get(featureType,'NA') == 'NA':
        strengths[target][featureType] = float(importance)
    else:
        strengths[target][featureType] = max(strengths[target][featureType],float(importance))
            
for target in strengths.keys():
    #print target,strengths[target].keys(),allFeatureTypes,'\n',
    
    totalStrength = 0.0
    for featureType in strengths[target].keys():
        totalStrength += strengths[target][featureType]
    strengths[target]['SUM'] = totalStrength


allFeatureTypes = list(allFeatureTypes)
allFeatureTypes.append('SUM')

outputList = []

for target in strengths.keys():

    strength = strengths[target]
    newLine = [target]
    
    for featureType in allFeatureTypes:
        newLine.append( strength.get(featureType,'NA') )

    outputList.append(newLine)

outputList = sorted(outputList, key=itemgetter(len(allFeatureTypes)), reverse=True)

allFeatureTypes.insert(0,'FEATURE')

listwriter.writerow(allFeatureTypes)

for line in outputList:
    listwriter.writerow(line)
    
